package client

import (
	"net/http"
	"time"
)

type RetryOption struct {
	maxRetries int
	backoff    time.Duration
}

func WithRetry(maxRetries int) HttpOption {
	return func(c *http.Client) {
		c.Transport = &retryTransport{
			transport: c.Transport,
			option: RetryOption{
				maxRetries: maxRetries,
				backoff:    time.Second,
			},
		}
	}
}

type retryTransport struct {
	transport http.RoundTripper
	option    RetryOption
}

func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	var resp *http.Response
	var err error

	for i := 0; i <= t.option.maxRetries; i++ {
		if i > 0 {
			time.Sleep(t.option.backoff)
		}

		if t.transport == nil {
			t.transport = http.DefaultTransport
		}

		resp, err = t.transport.RoundTrip(req)
		if err == nil && resp.StatusCode < 500 {
			return resp, nil
		}

		if resp != nil {
			resp.Body.Close()
		}
	}

	return resp, err
}
